import pyb
import sensor, image, time, math
import os, nncu
from machine import UART
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QVGA)
sensor.set_brightness(1000)
sensor.skip_frames(time = 1000)
sensor.set_auto_gain(False)
sensor.set_auto_whitebal(False)
clock = time.clock()
uart = UART(1, baudrate=115200)
threshold = [(42, 62, 35, 81, -3, 53)]
rect_search=(110,30,110,120)
net_path = "model.nncu"
labels = [line.rstrip() for line in open("/sd/labels_number.txt")]
net = nncu.load(net_path, load_to_fb=True)
num_list=[0,0,0,0,0,0,0,0,0]
index=0
dest=0
def send_inf(cmd_back):
    temp=[22,48,cmd_back]
    for i in range(20):
        for j in temp:
            uart.writechar(j)
            time.sleep(10)
def find_max(blobs):
    max_size=0
    for blob in blobs:
        if blob.pixels() > max_size:
            max_blob=blob
            max_size = blob.pixels()
    return max_blob
flag=0
while True:
    clock.tick()
    if uart.any():
        if uart.read(1) == b'1':
            if uart.read(1) == b'2':
                ttt=uart.read(1)
                if  ttt== b'3':
                    flag=1
                elif ttt == b'4':
                    flag=2
    print(flag)
    if flag==1:
        img = sensor.snapshot()
        img.draw_rectangle(rect_search, color = (0, 255, 0))
        for r in img.find_rects(threshold = 50000,roi=rect_search):
            img.draw_rectangle(r.rect(), color = (255, 0, 0))
            img1 = img.copy(r.rect())
            for obj in nncu.classify(net , img1, min_scale=1.0, scale_mul=0.5, x_overlap=0.0, y_overlap=0.0):
                sorted_list = sorted(zip(labels, obj.output()), key = lambda x: x[1], reverse = True)
                num=int(sorted_list[0][0])
                if num==0:
                    num=1
                if num==9:
                    num=2
                img.draw_string(r.x(),r.y(),str(num),color=(0,0,255),scale=3)
                index+=1
                num_list[num]+=1
                if index==6:
                    temp=sorted(zip(labels,num_list),key=lambda x:x[1],reverse=True)
                    if temp[0][1]>=4 and dest==0:
                        dest=int(temp[0][0])
                        send_inf(dest)
                        print(dest)
                        flag=0
                    index=0
                    num_list=[0,0,0,0,0,0,0,0,0]
    if flag==2:
        cmd_back=0
        index=0
        num_list=[0,0,0,0,0,0,0,0,0]
        while not cmd_back:
            img = sensor.snapshot()
            blobs = img.find_blobs(threshold,roi = (40,80,200,50))
            if blobs:
                maxb=find_max(blobs)
                img.draw_rectangle(maxb.rect())
                x_min=maxb.cx()+10
            else:
                x_min=60
            img.draw_rectangle((x_min,30,320,140),color=(0,255,0))
            for r in img.find_rects(threshold = 50000,roi=(x_min,30,320,140)):
                img.draw_rectangle(r.rect(), color = (255, 0, 0))
                img1 = img.copy(r.rect())
                for obj in nncu.classify(net , img1, min_scale=1.0, scale_mul=0.5, x_overlap=0.0, y_overlap=0.0):
                    sorted_list = sorted(zip(labels, obj.output()), key = lambda x: x[1], reverse = True)
                    num=int(sorted_list[0][0])
                    if num==0:
                        num=1
                    if num==9:
                        num=2
                    img.draw_string(r.x(),r.y(),str(num),color=(0,0,255),scale=3)
                    index+=1
                    num_list[int(num)]+=1
                    if index==8:
                        temp=sorted(zip(labels,num_list),key=lambda x:x[1],reverse=True)
                        if (temp[0][1] and int(temp[0][0])==dest) or (temp[1][1] and int(temp[1][0])==dest):
                            cmd_back=2
                        else:
                            cmd_back=1
                        send_inf(cmd_back)
                        flag=0
                        print(cmd_back)
                        index=0
                        num_list=[0,0,0,0,0,0,0,0,0]
